import py_trees
import time


class CheckBattery(py_trees.behaviour.Behaviour):
    def __init__(self, name):
        super().__init__(name)
        self.battery_level = 80  # 模拟电量

    def update(self):
        if self.battery_level > 30:
            return py_trees.common.Status.SUCCESS
        else:
            return py_trees.common.Status.FAILURE


class IsTargetVisible(py_trees.behaviour.Behaviour):
    def update(self):
        # 模拟目标检测
        target_visible = True  # 改为False测试不同行为
        return (
            py_trees.common.Status.SUCCESS
            if target_visible
            else py_trees.common.Status.FAILURE
        )


class MoveTo(py_trees.behaviour.Behaviour):
    def __init__(self, name, target):
        super().__init__(name)
        self.target = target

    def update(self):
        print(f"移动到 {self.target}")
        return py_trees.common.Status.SUCCESS


class GrabObject(py_trees.behaviour.Behaviour):
    def __init__(self, name):
        super().__init__(name)

    def update(self):
        print("抓取物体")
        return py_trees.common.Status.SUCCESS


class PlaceObject(py_trees.behaviour.Behaviour):
    def __init__(self, name, location):
        super().__init__(name)
        self.location = location

    def update(self):
        print(f"放置物体到 {self.location}")
        return py_trees.common.Status.SUCCESS


# 创建行为树
def create_behavior_tree():
    root = py_trees.composites.Sequence("根序列", memory=True)

    # 检查条件
    check_battery = CheckBattery("检查电量")
    check_visibility = IsTargetVisible("检查目标可见性")

    # 移动和操作
    move_to_object = MoveTo("移动到物体", "桌子A")
    grab = GrabObject("抓取物体")
    move_to_destination = MoveTo("移动到目的地", "架子B")
    place = PlaceObject("放置物体", "架子B")

    # 构建树结构
    root.add_children(
        [
            check_battery,
            py_trees.composites.Sequence("抓取序列", memory=True).add_children(
                [check_visibility, move_to_object, grab]
            ),
            py_trees.composites.Sequence("放置序列", memory=True).add_children(
                [move_to_destination, place]
            ),
        ]
    )

    return root


# 执行行为树
def main():
    behaviour_tree = create_behavior_tree()

    # 设置树
    behaviour_tree.setup(timeout=15)
    # 打印树结构
    print("行为树结构:")
    print(py_trees.display.ascii_tree(behaviour_tree))

    # 执行树
    print("=== 开始执行行为树 ===")
    for i in range(5):
        print(f"\n--- 第 {i+1} 次tick ---")
        behaviour_tree.tick_once()
        time.sleep(1)

    print("\n=== 行为树执行完成 ===")


if __name__ == "__main__":
    main()
